"""Model management service for PolyModel Agent"""

from typing import Dict, Type, Optional, Any
import asyncio
from pathlib import Path
import logging

from modules.base.base_model import BaseModel
from modules.base.exceptions import ModelNotFoundError, ModelLoadError
from config.settings import Config

class ModelManager:
    """Manager class for handling AI models"""
    
    def __init__(self):
        self.config = Config()
        self.logger = logging.getLogger(self.__class__.__name__)
        self._models: Dict[str, BaseModel] = {}
        self._model_classes: Dict[str, Type[BaseModel]] = {}
        
    def register_model(self, model_type: str, model_class: Type[BaseModel]) -> None:
        """
        Register a new model class.
        
        Args:
            model_type (str): Type identifier for the model
            model_class (Type[BaseModel]): Model class to register
        """
        self._model_classes[model_type] = model_class
        self.logger.info(f"Registered model type: {model_type}")
    
    async def load_model(
        self, 
        model_type: str, 
        model_name: Optional[str] = None,
        config: Optional[Dict[str, Any]] = None
    ) -> BaseModel:
        """
        Load a model instance.
        
        Args:
            model_type (str): Type of model to load
            model_name (Optional[str]): Specific model name
            config (Optional[Dict]): Model configuration
            
        Returns:
            BaseModel: Loaded model instance
            
        Raises:
            ModelNotFoundError: If model type is not registered
            ModelLoadError: If model loading fails
        """
        if model_type not in self._model_classes:
            raise ModelNotFoundError(f"Model type not found: {model_type}")
            
        if not model_name:
            model_name = self.config.model_configs[model_type]["default"]
            
        model_key = f"{model_type}:{model_name}"
        
        if model_key not in self._models:
            try:
                model_class = self._model_classes[model_type]
                model = model_class(model_name, config)
                await model.load()
                self._models[model_key] = model
                self.logger.info(f"Loaded model: {model_key}")
            except Exception as e:
                raise ModelLoadError(f"Failed to load model {model_key}: {str(e)}")
                
        return self._models[model_key]
    
    async def unload_model(self, model_type: str, model_name: str) -> None:
        """
        Unload a model instance.
        
        Args:
            model_type (str): Type of model to unload
            model_name (str): Name of model to unload
        """
        model_key = f"{model_type}:{model_name}"
        if model_key in self._models:
            del self._models[model_key]
            self.logger.info(f"Unloaded model: {model_key}")
    
    def get_available_models(self, model_type: Optional[str] = None) -> Dict[str, list]:
        """
        Get available models.
        
        Args:
            model_type (Optional[str]): Specific model type to query
            
        Returns:
            Dict[str, list]: Available models by type
        """
        if model_type:
            return {
                model_type: self.config.model_configs[model_type]["available"]
            }
        return {
            k: v["available"] 
            for k, v in self.config.model_configs.items()
        }
